import numpy as np
import matplotlib.pyplot as plt
from scipy.interpolate import griddata

# 创建虚拟数据
def generate_2d_data():
    x = np.linspace(0, 4, 5)
    y = np.linspace(0, 4, 5)
    X, Y = np.meshgrid(x, y)
    Z = np.sin(X) * np.cos(Y)
    return X, Y, Z

def generate_3d_data():
    x = np.linspace(0, 4, 5)
    y = np.linspace(0, 4, 5)
    z = np.linspace(0, 4, 5)
    X, Y, Z = np.meshgrid(x, y, z)
    V = np.sin(X) * np.cos(Y) * np.sin(Z)
    return X, Y, Z, V

# 进行双线性插值
def bilinear_interpolation(X, Y, Z):
    xi = np.linspace(0, 4, 100)
    yi = np.linspace(0, 4, 100)
    XI, YI = np.meshgrid(xi, yi)
    Zi = griddata((X.flatten(), Y.flatten()), Z.flatten(), (XI, YI), method='linear')
    return XI, YI, Zi

# 进行三线性插值
def trilinear_interpolation(X, Y, Z, V):
    xi = np.linspace(0, 4, 20)
    yi = np.linspace(0, 4, 20)
    zi = np.linspace(0, 4, 20)
    XI, YI, ZI = np.meshgrid(xi, yi, zi)
    Vi = griddata((X.flatten(), Y.flatten(), Z.flatten()), V.flatten(), (XI, YI, ZI), method='linear')
    return XI, YI, ZI, Vi

# 绘制双线性插值结果
def plot_bilinear(X, Y, Z, XI, YI, Zi):
    plt.figure(figsize=(10, 8))
    plt.subplot(2, 2, 1)
    plt.pcolormesh(X, Y, Z, shading='auto', cmap='jet')
    plt.title("Original Data (2D)")
    plt.colorbar()

    plt.subplot(2, 2, 2)
    plt.contourf(XI, YI, Zi, levels=15, cmap='jet')
    plt.title("Bilinear Interpolated Data (2D)")
    plt.colorbar()

# 绘制三线性插值结果
def plot_trilinear(X, Y, Z, V, XI, YI, ZI, Vi):
    fig = plt.figure(figsize=(10, 8))

    # 绘制原始数据切片
    ax1 = fig.add_subplot(2, 2, 1, projection='3d')
    slice_z = V[:, :, 2]
    ax1.plot_surface(X[:, :, 2], Y[:, :, 2], slice_z, cmap='jet')
    ax1.set_title('Original Data (Slice z=2)')

    # 绘制插值后的切片
    ax2 = fig.add_subplot(2, 2, 2, projection='3d')
    slice_zi = Vi[:, :, 10]
    ax2.plot_surface(XI[:, :, 10], YI[:, :, 10], slice_zi, cmap='jet')
    ax2.set_title('Trilinear Interpolation (Slice z=10)')

    # 绘制体积可视化
    ax3 = fig.add_subplot(2, 2, 3, projection='3d')
    ax3.scatter(X.flatten(), Y.flatten(), Z.flatten(), c=V.flatten(), cmap='jet')
    ax3.set_title('Original 3D Data Points')

    ax4 = fig.add_subplot(2, 2, 4, projection='3d')
    ax4.scatter(XI.flatten(), YI.flatten(), ZI.flatten(), c=Vi.flatten(), cmap='jet')
    ax4.set_title('Trilinear Interpolated 3D Data')

# 主程序
X2, Y2, Z2 = generate_2d_data()
XI2, YI2, Zi2 = bilinear_interpolation(X2, Y2, Z2)
plot_bilinear(X2, Y2, Z2, XI2, YI2, Zi2)

X3, Y3, Z3, V3 = generate_3d_data()
XI3, YI3, ZI3, Vi3 = trilinear_interpolation(X3, Y3, Z3, V3)
plot_trilinear(X3, Y3, Z3, V3, XI3, YI3, ZI3, Vi3)

plt.tight_layout()
plt.show()